package hw2.nsl;

import hw2.sign.FileIO;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;


public class Alice {
	
	static PublicKey bobsKey;
	static PrivateKey alicePriv;
	static final int portno = 7000;
	static final int malportno = 7001;
	static boolean encryptMode = true;
	static Integer myNonce = 0;

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		System.out.println("\nYou are running the Alice program. (enter quit to exit program)");
		if (args.length > 0){
			if (args[0].equalsIgnoreCase("plain")){
				encryptMode = false;
				System.out.println("Encryption has been turned off.");
			}
		}
		BufferedReader in = new BufferedReader( new InputStreamReader(System.in ) );
        String s = "";
        new ListenerThread().start();
        
        readKeys();
		
        while( true ) {
        	try{
        		System.out.print("Enter a nonce(or r to generate a random nonce) >>");
        		s = in.readLine().toLowerCase();
        		if( "quit".equalsIgnoreCase( s ) )
        			break;
        		if ("r".equalsIgnoreCase(s)){
        			//generate a random nonce
        			SecureRandom sr = SecureRandom.getInstance("SHA1PRNG");
        			myNonce = sr.nextInt();
        		}
        		else{
        			myNonce = Integer.parseInt(s);
        		}
        		// create string to Enc/send to "Bob"
        		String strMessage = "Alice " + myNonce;
        		byte[] message = strMessage.getBytes();
        		if(encryptMode){
        			//Encrypt message if encryption is turned on
        			message = encryptToBob(message);
        		}
        		NetworkMessage msg = new NetworkMessage(message, "Alice");
        		sendMessage(msg);
        		
        		
        	}catch(Exception e){
        		e.printStackTrace();
        	}
        }
        System.out.println( "Goodbye." );
        System.exit(0);
	}
	
	/**
	 * Read in Bob's public key and alice's private key
	 */
	private static void readKeys() {
		try {
			byte[] b = FileIO.getBytesFromFile("alice_priv.bin");
			PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(b);
			alicePriv =  KeyFactory.getInstance( "RSA" ).generatePrivate( spec );

			b = FileIO.getBytesFromFile("bob_pub.bin");
			X509EncodedKeySpec spec2 = new X509EncodedKeySpec(b);
			bobsKey =  KeyFactory.getInstance( "RSA" ).generatePublic( spec2 );
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
	}

	public static void sendMessage(NetworkMessage m) throws Exception{
		Socket clientSocket = new Socket("127.0.0.1",malportno);
//		Socket clientSocket = new Socket("127.0.0.1",7002);
		ObjectOutputStream out = new ObjectOutputStream(clientSocket.getOutputStream());
		System.out.println("About to send: " + m.toHexString());
		out.writeObject(m);
		
		out.flush();
		out.close();
	}
	
	/**
	 * Enc with bob's public key and get the enc bytes back
	 * @param message
	 * @return
	 * @throws NoSuchAlgorithmException
	 * @throws NoSuchPaddingException
	 * @throws InvalidKeyException
	 * @throws IllegalBlockSizeException
	 * @throws BadPaddingException
	 */
	private static byte[] encryptToBob(byte[] message) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
		Cipher enc = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-256ANDMGF1PADDING");
		enc.init(Cipher.ENCRYPT_MODE, bobsKey);
		byte[] encMsg = enc.doFinal(message);
		return encMsg;
	}
	
	

    /**
	 * Decrypt our message
	 * @param message - bytes which were enced with our (Alice's) public key
	 * @return
	 * @throws NoSuchPaddingException 
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeyException 
	 * @throws BadPaddingException 
	 * @throws IllegalBlockSizeException 
	 */
	private static byte[] decrypt(byte[] message) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
		Cipher dec = Cipher.getInstance("RSA/ECB/OAEPWITHSHA-256ANDMGF1PADDING");
		dec.init(Cipher.DECRYPT_MODE, alicePriv);
		return dec.doFinal(message);
	}

	
	
	public static void handleMessage(NetworkMessage msg) throws Exception {
		byte[] text = msg.text;
		if (encryptMode){
			try{
				text = decrypt(text);
			}catch(Exception up){
				System.err.println("\nERROR: Failed to decrypt message");
				return;
			}
		}
		NetworkMessage m = new NetworkMessage(text, "Alice");
		String[] vals = m.parseFields();
		if (vals.length != 3){
			System.err.println("\nERROR: Recived a message without 3 values");
			return;
		}
		//check from Bob
		if (!vals[0].equalsIgnoreCase("Bob")){
			System.err.println("\nERROR: Received response from " + vals[0] + " but expected Bob.");
		}
		//check nonce
		else if(Integer.parseInt(vals[1]) != myNonce){
			System.err.println("ERROR: Nonce didn't match, " + vals[1] + " != " + myNonce);
		}
		// create message to send to Bob (step 3 of protocol)
		else{
			byte[] bobNonce = vals[2].getBytes();
			if (encryptMode){
				bobNonce = encryptToBob(bobNonce);
			}
			NetworkMessage toSend = new NetworkMessage(bobNonce, "Alice");
			Alice.sendMessage(toSend);
			System.out.println("\nSUCCESS: I accept that I am talking to Bob.");
		}
	}
}

class ListenerThread extends Thread{
	ServerSocket sock;
	
	@Override
	public void run() {
		try {
			sock = new ServerSocket(Alice.portno);
		} catch (IOException e) {
			e.printStackTrace();
		}
		do{
			try{
				Socket client = sock.accept();
				ObjectInputStream in = new ObjectInputStream(client.getInputStream());
				
				NetworkMessage msg = (NetworkMessage)in.readObject();		
				System.out.println("Just received: " + msg.toHexString());
				client.close();
				
				Alice.handleMessage(msg);
			}catch(Exception e){
				e.printStackTrace();
			}				
		}while(true);
	}		
}
